package com.cc.flink.sink;

import com.cc.flink.config.JobConfig;
import com.cc.flink.exception.JobExceptionHandler;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 自定义Doris批量写入Sink
 * 实现Flink数据到Doris数据库的高效批量写入
 * <p>
 * 核心特性：
 * 1. 支持精确一次语义（Exactly-Once），通过检查点机制保证
 * 2. 批量写入优化，减少数据库连接开销
 * 3. 错误重试机制，提高写入可靠性
 * 4. 支持事务，确保数据写入一致性
 */
public class DorisBulkJdbcSink extends RichSinkFunction<Tuple2<String, String>> implements CheckpointedFunction {
    
    /**
     * 日志记录器，用于记录写入状态和错误信息
     */
    private static final Logger logger = LoggerFactory.getLogger(DorisBulkJdbcSink.class);
    
    // 配置相关成员变量
    private final JobConfig config;         // 作业配置对象，包含数据库连接信息等
    private final int batchSize;            // 批量写入的最大记录数
    private final long flushIntervalMs;     // 定时刷新间隔（毫秒）
    private final int maxRetries;           // 写入失败后的最大重试次数
    
    // 数据库相关成员变量
    private transient Connection connection;           // 数据库连接（transient标记非序列化字段）
    private transient PreparedStatement preparedStatement; // 预编译SQL语句
    
    // 批处理相关成员变量
    private final List<Tuple2<String, String>> batchBuffer; // 批量数据缓冲区
    private transient long lastFlushTime;                   // 上次刷新时间戳
    
    /**
     * 用于保存检查点状态的列表，确保故障恢复时不丢失数据
     * 该状态会参与Flink的检查点机制
     */
    private transient ListState<Tuple2<String, String>> checkpointedState;
    
    /**
     * 构造函数，使用完全自定义参数
     * @param config 作业配置对象，包含数据库连接信息
     * @param batchSize 批处理大小，达到此数量时触发写入
     * @param flushIntervalMs 刷新间隔（毫秒），达到此时间时触发写入
     * @param maxRetries 最大重试次数，提高写入可靠性
     */
    public DorisBulkJdbcSink(JobConfig config, int batchSize, long flushIntervalMs, int maxRetries) {
        this.config = config;
        this.batchSize = batchSize;
        this.flushIntervalMs = flushIntervalMs;
        this.maxRetries = maxRetries;
        this.batchBuffer = new ArrayList<>(batchSize); // 预分配容量，减少动态扩容开销
    }
    
    /**
     * 构造函数 - 使用配置中的默认值
     * 提供更简化的创建方式，适用于大多数场景
     * @param config 作业配置对象
     */
    public DorisBulkJdbcSink(JobConfig config) {
        // 默认值：使用配置中的批处理大小，1000ms刷新间隔，3次重试
        this(config, (int) config.getMaxRecordsPerBatch(), 1000, 3);
    }
    
    /**
     * 初始化Sink实例
     * 当Flink任务启动时调用，用于初始化数据库连接和资源
     * @param parameters Flink配置参数
     * @throws Exception 初始化过程中发生的异常
     */
    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        // 初始化最后刷新时间，用于定时刷新机制
        this.lastFlushTime = System.currentTimeMillis();
        // 初始化数据库连接和预编译语句
        initConnection();
    }
    
    /**
     * 处理单条输入记录
     * 当有新数据到达时调用，实现批量收集和刷新逻辑
     * @param value 输入数据，Tuple2<key, value>格式
     * @param context Flink上下文，可获取任务状态等信息
     * @throws Exception 数据处理过程中发生的异常
     */
    @Override
    public void invoke(Tuple2<String, String> value, Context context) throws Exception {
        synchronized (batchBuffer) {
            // 添加到批处理缓冲区
            batchBuffer.add(value);
            
            // 触发条件：数据量达到阈值 或 时间间隔达到阈值
            long currentTime = System.currentTimeMillis();
            if (batchBuffer.size() >= batchSize || (currentTime - lastFlushTime) >= flushIntervalMs) {
                // 满足任一条件则执行批量写入
                flushBatch();
                // 更新最后刷新时间
                lastFlushTime = currentTime;
            }
        }
    }
    
    /**
     * 关闭资源
     * 当Flink任务停止时调用，确保资源正确释放和数据完整性
     * @throws Exception 资源关闭过程中发生的异常
     */
    @Override
    public void close() throws Exception {
        // 关闭前必须刷新剩余数据，避免数据丢失
        if (!batchBuffer.isEmpty()) {
            try {
                flushBatch();
            } catch (Exception e) {
                logger.error("关闭时刷新剩余数据失败: {}", e.getMessage(), e);
            }
        }
        
        // 关闭数据库资源，按顺序关闭：先关闭语句，再关闭连接
        if (preparedStatement != null) {
            try {
                preparedStatement.close();
            } catch (SQLException e) {
                logger.warn("关闭预编译语句失败: {}", e.getMessage(), e);
            }
        }
        
        if (connection != null) {
            try {
                connection.close();
            } catch (SQLException e) {
                logger.warn("关闭数据库连接失败: {}", e.getMessage(), e);
            }
        }
        
        super.close();
    }
    
    /**
     * 初始化数据库连接和预编译语句
     * 获取配置信息并建立到Doris数据库的连接
     * @throws SQLException 数据库连接失败时抛出
     */
    private void initConnection() throws SQLException {
        try {
            // 从配置获取Doris数据库连接信息
            String url = config.getDorisUrl();
            String username = config.getDorisUsername();
            String password = config.getDorisPassword();
            
            // 建立数据库连接
            connection = DriverManager.getConnection(url, username, password);
            
            // 禁用自动提交，启用事务控制，保证数据一致性
            connection.setAutoCommit(false);
            
            // 创建预编译语句，提高执行效率和安全性
            preparedStatement = connection.prepareStatement(config.getDorisInsertSql());
            
            logger.info("Doris数据库连接初始化成功");
        } catch (SQLException e) {
            logger.error("初始化Doris数据库连接失败", e);
            throw e;
        }
    }
    
    /**
     * 批量写入数据到Doris数据库
     * 实现了重试机制和错误处理，确保数据写入可靠性
     * @throws Exception 批量写入失败且重试次数用尽时抛出
     */
    private void flushBatch() throws Exception {
        // 如果缓冲区为空，直接返回
        if (batchBuffer.isEmpty()) {
            return;
        }
        
        logger.info("将{}条记录批量写入Doris", batchBuffer.size());
        
        // 执行批量写入，支持重试
        int retryCount = 0;
        boolean success = false;
        Exception lastException = null;
        
        while (!success && retryCount <= maxRetries) {
            try {
                // 重试时重新建立连接，确保连接有效
                if (retryCount > 0) {
                    logger.info("批量写入重试, 尝试 {}/{}", retryCount, maxRetries);
                    initConnection();
                    // 指数退避策略，避免频繁重试
                    long backoffTime = (long) (Math.pow(2, retryCount) * 1000);
                    TimeUnit.MILLISECONDS.sleep(backoffTime);
                }
                
                // 批量插入数据
                for (Tuple2<String, String> record : batchBuffer) {
                    preparedStatement.setString(1, record.f0); // 设置第一个参数（key）
                    preparedStatement.setString(2, record.f1); // 设置第二个参数（value）
                    preparedStatement.addBatch(); // 添加到批处理
                }
                
                // 执行批处理并提交事务
                int[] results = preparedStatement.executeBatch();
                connection.commit();
                
                logger.info("成功写入{}条记录到Doris", results.length);
                
                // 清空缓冲区，准备接收新数据
                batchBuffer.clear();
                success = true;
            } catch (Exception e) {
                lastException = e;
                retryCount++;
                
                // 发生异常时必须回滚事务，防止数据不一致
                try {
                    if (connection != null) {
                        connection.rollback();
                    }
                } catch (SQLException rollbackEx) {
                    logger.warn("事务回滚失败: {}", rollbackEx.getMessage(), rollbackEx);
                }
                
                // 判断是否应该继续重试
                if (retryCount > maxRetries || !JobExceptionHandler.handleDatabaseException(e, batchBuffer)) {
                    logger.error("达到最大重试次数{}后写入失败: {}", maxRetries, e.getMessage(), e);
                    // 记录失败的数据到日志，便于后续排查
                    for (Tuple2<String, String> failedRecord : batchBuffer) {
                        logger.error("写入失败的记录: {}", failedRecord);
                    }
                    // 清空缓冲区，避免重复处理
                    batchBuffer.clear();
                    throw e;
                }
            }
        }
    }
    
    /**
     * 初始化检查点状态
     * 当Flink任务启动或从检查点恢复时调用，用于设置状态存储和恢复数据
     * @param context 函数初始化上下文，提供状态存储访问
     * @throws Exception 状态初始化过程中发生的异常
     */
    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        // 创建列表状态描述符，定义状态的名称和类型
        // 这里使用ListState存储待处理的数据，确保数据可恢复
        ListStateDescriptor<Tuple2<String, String>> descriptor = 
                new ListStateDescriptor<>(
                        "doris-batch-buffer", // 状态名称，唯一标识
                        Types.TUPLE(Types.STRING, Types.STRING) // 状态类型：Tuple2<String, String>
                );
        
        // 从操作符状态存储中获取列表状态
        // 使用OperatorStateStore而不是KeyedStateStore，因为Sink通常不关心键分组
        checkpointedState = context.getOperatorStateStore().getListState(descriptor);
        
        // 检查是否是从检查点恢复（任务重启）
        if (context.isRestored()) {
            logger.info("从检查点恢复状态");
            // 遍历所有保存的状态数据，重新添加到批处理缓冲区
            for (Tuple2<String, String> value : checkpointedState.get()) {
                batchBuffer.add(value);
            }
            logger.info("从检查点成功恢复{}条记录", batchBuffer.size());
            // 恢复后立即尝试写入，避免数据积压
            if (!batchBuffer.isEmpty()) {
                flushBatch();
            }
        }
    }
    
    /**
     * 保存状态到检查点
     * 当Flink触发检查点时调用，确保数据的精确一次语义
     * @param context 函数快照上下文，提供检查点元数据
     * @throws Exception 状态保存过程中发生的异常
     */
    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        // 在保存新状态前，先清除之前的状态数据
        checkpointedState.clear();
        
        // 将当前批处理缓冲区中的所有数据添加到检查点状态中
        // 这些数据将在任务失败重启时被恢复
        for (Tuple2<String, String> value : batchBuffer) {
            checkpointedState.add(value);
        }
        
        logger.debug("快照状态：当前保存{}条记录", batchBuffer.size());
        // 注意：此方法在检查点触发时调用，此时数据尚未写入数据库
        // 只有当检查点确认完成后，数据才会被视为安全写入
    }
}